/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.util.*;

import edu.berkeley.nlp.util.Pair;
import edu.umd.clip.lm.util.BinaryPrefix;
import edu.umd.clip.lm.util.ProbMath;
import edu.umd.clip.lm.util.SimpleList;
import edu.umd.clip.lm.factors.Dictionary;
import edu.umd.clip.lm.factors.FactorTuple;
import edu.umd.clip.lm.model.Experiment;
import edu.umd.clip.lm.model.ProbTree;
import edu.umd.clip.lm.model.decoding.VariationalTaggerTrellisSlice.Trie;
import edu.umd.clip.lm.model.decoding.TrellisSlice.SliceItem;
/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class VariationalTagger extends Tagger {
	final int taggerCtxSize;
	
	public static class VariationalPath implements Comparable<VariationalPath>, TaggedPath {
		final double logProb;
		final SimpleList<Integer> tags;
		private static byte tagFactor;
		
		public VariationalPath(double logProb, SimpleList<Integer> tags) {
			this.logProb = logProb;
			this.tags = tags;
		}

		public VariationalPath append(int tag, double logProb) {
			VariationalPath variationalPath = new VariationalPath(this.logProb + logProb, tags.add(tag));
			return variationalPath;
		}
		
		@Override
		public int compareTo(VariationalPath other) {
			return (int) Math.signum(logProb - other.logProb);
		}
		
		public int[] getContext(int ctxSize) {
			int ctx[] = new int[ctxSize];
			SimpleList<Integer> lst = tags;
			for(int i=0; i<ctxSize; ++i) {
				if (lst.head() != null) {
					ctx[i] = lst.head().intValue();
					lst = lst.tail();
				} else {
					ctx[i] = Dictionary.getStart();
				}
			}
			return ctx;
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.decoding.TaggedPath#getLogProb()
		 */
		public double getLogProb() {
			return logProb;
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.decoding.TaggedPath#getOrderedTags()
		 */
		public int[] getOrderedTags() {
			int t[] = new int[tags.length()];
			SimpleList<Integer> tail = tags;
			for(int i=1; i<=t.length; ++i) {
				// reverse to the proper order
				t[t.length-i] = tail.head().intValue();
				tail = tail.tail();
			}
			return t;
		}
		
		public String toString() {
			StringBuilder sb = new StringBuilder();
			Dictionary dict = Experiment.getInstance().getTupleDescription().getDictionary(tagFactor);
			
			int tags[] = getOrderedTags();
			for(int tag : tags) {
				sb.append(dict.getWord(tag));
				sb.append(' ');
			}
			sb.append('(');
			sb.append(logProb);
			sb.append(')');
			return sb.toString();
		}
	}
	
	public VariationalTagger(Decoder decoder, int beamSize, String tagFactor) {
		super(decoder, beamSize, tagFactor);
		VariationalPath.tagFactor = tagFactorIdx;
		taggerCtxSize = decoder.getLm().getOrder() - 1;
	}
	
	public Pair<SortedSet<? extends TaggedPath>,Double> evaluateSentence(Collection<FactorTuple> sentence) {
		SentenceContext ctx = decoder.createContext();
		VariationalTaggerContext variationalTaggerContext = new VariationalTaggerContext(tags);
		final int lmCtxSize = decoder.getLm().getOrder() - 1;
		//HashMap<TaggerContext.Context,List<ViterbiPath>> currentPathsHash = taggerContext.currentPathsHash;
		//TreeSet<ViterbiPath> currentPaths = taggerContext.currentPaths;
		
		VariationalPath initPath;
		{
			SimpleList<Integer> seq = new SimpleList<Integer>();
			for(int i=0; i<lmCtxSize; ++i) {
				seq = seq.add(edu.umd.clip.lm.factors.Dictionary.getStart());
			}
			initPath = new VariationalPath(0.0, seq);
		}
		variationalTaggerContext.currentPaths.add(initPath);
		
		{
			List<VariationalPath> lst = new LinkedList<VariationalPath>();
			lst.add(initPath);
			variationalTaggerContext.currentPathsHash.put(new VariationalTaggerContext.Context(initPath.getContext(lmCtxSize)), lst);
		}
		
		long overtMask = Experiment.getInstance().getTupleDescription().getOvertFactorsMask();
		for(FactorTuple tuple : sentence) {
			tuple = new FactorTuple(tuple.getBits() & overtMask);
			ctx = ctx.addWord(tuple.getBits());
			decoder.partition(ctx);
			if (decoder.getDebug() > 0) {
				decoder.debugContext(ctx);
				if (decoder.getDebug() > 1) {
					ctx.getTrellis().dumpTrellis(String.format("trellis-tagger-%d-%s.dot", 
							ctx.getTrellis().getCurrentSlice().getTime(), tuple.toStringNoNull().replace('/', ' ')));
				}
			}
			if (ctx.getCurrentSlice().getTime() <= lmCtxSize + 1) continue;
			
			long start = System.nanoTime();
			
			VariationalTaggerTrellisSlice slice = makeTaggerSlice(variationalTaggerContext, ctx.getCurrentSlice().getPrevious());
			
			long end = System.nanoTime();
			long updatePathsDelta = end - start;
			
			updatePathsTime += updatePathsDelta;
			
			start = end;
			
			updatePaths(variationalTaggerContext, slice);
			
			end = System.nanoTime();
			long makingTaggerSliceDelta = end - start;
			
			makingTaggerSliceTime += makingTaggerSliceDelta;
			
			if (decoder.getDebug() > 0) {
				System.out.printf("tagger: %d paths, %d items in tag hash, update: %fs, slice: %fs\n", 
						variationalTaggerContext.currentPaths.size(), variationalTaggerContext.itemTagHash.size(),
						(double)updatePathsDelta/1000000000,
						(double)makingTaggerSliceDelta/1000000000);
				printSlice(slice);
			}
		}
		// do the sentence terminator
		long start = System.nanoTime();

		VariationalTaggerTrellisSlice slice = makeTaggerSlice(variationalTaggerContext, ctx.getCurrentSlice());
		
		long end = System.nanoTime();
		updatePathsTime += end - start;
		start = end;
		
		updatePaths(variationalTaggerContext, slice);
		
		end = System.nanoTime();
		makingTaggerSliceTime += end - start;

		if (decoder.getDebug() > 0) {
			printSlice(slice);
		}
		return new Pair<SortedSet<? extends TaggedPath>,Double>(variationalTaggerContext.currentPaths, ctx.getLogProb());
	}
	
	private void updatePaths(VariationalTaggerContext variationalTaggerContext, VariationalTaggerTrellisSlice slice) {
		//final int ctxSize = decoder.getLm().getOrder() - 1;

		TreeSet<VariationalPath> variationalPaths = new TreeSet<VariationalPath>();
		//HashMap<TaggerContext.Context, List<ViterbiPath>> pathsHash = new HashMap<TaggerContext.Context, List<ViterbiPath>>(taggerContext.currentPathsHash.size());
		
		double minProb = Double.NEGATIVE_INFINITY;
		//for(Map.Entry<TaggerContext.Context, List<ViterbiPath>> entry : taggerContext.currentPathsHash.entrySet()) {
		for(VariationalPath prevPath : variationalTaggerContext.currentPaths.descendingSet()) {
			VariationalTaggerContext.Context context = new VariationalTaggerContext.Context(prevPath.getContext(taggerCtxSize));

			if (variationalPaths.size() >= beamSize && prevPath.logProb < minProb) {
				// pruned
				break;
			}

			Trie trie = slice.findTrie(context.ctx);
			if (trie == null) {
				//System.out.printf("Tagger: no continuation found for %s\n", Arrays.toString(context.ctx));
				continue;
			}
			
			double ctxLogProb = Math.log10(trie.prob);
			for(Trie subTrie : trie.subTries) {
				double logProb = Math.log10(subTrie.prob) - ctxLogProb;
				if (variationalPaths.size() >= beamSize && logProb + prevPath.logProb < minProb) {
					// pruned
					continue;
				}
				VariationalPath newPath = prevPath.append(subTrie.tag, logProb);
				variationalPaths.add(newPath);
				if (variationalPaths.size() > beamSize) {
					variationalPaths.pollFirst();
					minProb = variationalPaths.first().logProb;
				}

				/*
				TaggerContext.Context newContext = new TaggerContext.Context(newPath.getContext(ctxSize));
				List<ViterbiPath> newPaths = pathsHash.get(newContext);
				if (newPaths == null) {
					newPaths = new LinkedList<ViterbiPath>();
					pathsHash.put(newContext, newPaths);							
				}
				newPaths.add(newPath);
				*/
			}
		}
		variationalTaggerContext.currentPaths = variationalPaths;
		//taggerContext.currentPathsHash = pathsHash;
	}
	
	private VariationalTaggerTrellisSlice makeTaggerSlice(VariationalTaggerContext context, TrellisSlice slice) {
		VariationalTaggerTrellisSlice taggerSlice = new VariationalTaggerTrellisSlice();
		
		// avoid considering new elements
		SliceItem elements[] = slice.getElements().toArray(new SliceItem[slice.getElements().size()]);
		
		for(SliceItem sliceItem : elements) {
			BinaryPrefix prefix = sliceItem.tree.getPrefix();
			for(int tag = 0; tag < tags.length; ++tag) {
				BinaryPrefix tagPrefix = tags[tag];
				if (tagPrefix == null) continue;
				
				if (tagPrefix.isPrefixOf(prefix)) {
					context.itemTagHash.put(sliceItem, tag);
					break;
				}
				if (prefix.isPrefixOf(tagPrefix)) {
					//ProbTree tmpBackup = (ProbTree) sliceItem.tree.clone();
					ProbTree probTree = sliceItem.tree.cut(tagPrefix);
					if (probTree == sliceItem.tree) {
						sliceItem.tree = probTree;
						context.itemTagHash.put(sliceItem, tag);
						break; // no need to check other tags
					}
					if (probTree != null) {
/*						
						System.out.printf("Tree prob was: %g, remainer: %g, new: %g\n", tmpBackup.getTreeProb(), 
								sliceItem.tree.getProbTree().getTreeProb(), probTree.getTreeProb());
						if (sliceItem.tree.getProbTree().getTreeProb() < ProbMath.MIN_FLOAT_PRECISION) {
							ProbTree test = tmpBackup.cut(tagPrefix);
						}
*/						
						SliceItem newItem = new SliceItem(probTree, sliceItem.getPrevProb(), (ClusterId)sliceItem.clusterid.clone());
						slice.split(sliceItem, newItem);
						//System.out.printf("tagger: cutting %s off %s\n", newItem, sliceItem);
						context.itemTagHash.put(newItem, tag);
					}
				}
			}
		}
		
		//int ctxLen = decoder.getLm().getOrder();
		//int ctx[] = new int[ctxLen];
		
		if (taggerCtxSize > 0) {
			TrellisSlice prevSlice = slice.getPrevious();
			for(SliceItem item : prevSlice.getElements()) {
				if (slice.prev2curr.containsKey(item)) {
					int tag = context.getItemTag(item);
					if (tag == -1) continue;
					Trie trie = taggerSlice.trie.addProb(tag, item.tree.getTotalProb());
					iterate(item, context, taggerSlice, trie, taggerCtxSize-1, item, item.tree.getTotalProb());
				}
			}
		} else {
			// special case for unigram tagger should it ever be needed
			for(SliceItem item : slice.getElements()) {
				int tag = context.getItemTag(item);
				taggerSlice.trie.addProb(tag, item.tree.getTotalProb());
			}
		}
		
		/*
		for(SliceItem item : slice.getElements()) {
			int tag = context.getItemTag(item);
			if (tag != -1) {
				ctx[ctxLen-1] = tag;
				iterate(context, taggerSlice, ctx, 0, item, item.tree.getTotalProb());
			}
		}
		*/
		// clean the tag hash
		{
			TrellisSlice s = slice;
			for(int i=1; i<=taggerCtxSize; ++i) {
				if (s.getTime() <= 1) break;
				s = s.getPrevious();
			}
			if (s != null) {
				for(SliceItem item : s.getElements()) {
					context.itemTagHash.remove(item);
				}
			}
		}
		return taggerSlice;
	}
	
	private void iterate(final SliceItem frontItem, final VariationalTaggerContext context, 
			final VariationalTaggerTrellisSlice taggerSlice, Trie trie, int stepsLeft, final SliceItem item, double prob) 
	{
		if (stepsLeft == 0) {
			final TrellisSlice slice = frontItem.slice.getNext();
			Collection<SliceItem> nextItems = slice.prev2curr.get(frontItem);
			for(SliceItem nextItem : nextItems) {
				int tag = context.getItemTag(nextItem);
				if (tag == -1) continue;
				trie.addProb(tag, prob * nextItem.tree.getTotalProb());
			}
			return;
		}
		
		Collection<SliceItem> prevs = null;
		
		if (item != null) {
			prevs = item.slice.getPreviousItems(item);
		}
		
		if (prevs == null) {
			int tag = Dictionary.getStart();
			Trie subTrie = trie.addProb(tag, prob);
			iterate(frontItem, context, taggerSlice, subTrie, stepsLeft-1, null, prob);
			return;
		}
		
		for(SliceItem prevItem : prevs) {
			int tag = context.getItemTag(prevItem);
			if (tag == -1) continue;
			Trie subTrie = trie.addProb(tag, prob * prevItem.tree.getTotalProb());
			iterate(frontItem, context, taggerSlice, subTrie, stepsLeft-1, prevItem, prob * prevItem.tree.getTotalProb());
		}
	}

	private void printSlice(VariationalTaggerTrellisSlice slice) {
		int ctx[] = new int[taggerCtxSize];
		if (slice.trie.subTries != null) {
			for(Trie subTrie: slice.trie.subTries) {
				ctx[0] = subTrie.tag;
				printSliceIter(subTrie, taggerCtxSize, ctx);
			}
		}
	}

	private void printSliceIter(Trie trie, int stepsLeft, int ctx[]) {
		if (stepsLeft > 1) {
			if (trie.subTries != null) {
				for(Trie subTrie: trie.subTries) {
					ctx[stepsLeft-1] = subTrie.tag;
					printSliceIter(subTrie, stepsLeft-1, ctx);
				}
			}
		} else {
			Dictionary dict = Experiment.getInstance().getTupleDescription().getDictionary(VariationalPath.tagFactor);

			for(int i=ctx.length-1; i>=0; --i) {
				System.out.print(dict.getWord(ctx[i]));
				System.out.print(' ');
			}
			System.out.printf("%f", Math.log10(trie.prob));
			System.out.print(" -> ( ");
			if (trie.subTries != null) {
				for(Trie subTrie : trie.subTries) {
					System.out.printf("%s:%f ", dict.getWord(subTrie.tag), Math.log10(subTrie.prob));
				}
			}
			System.out.print(")\n");
		}
	}
}
